# %%
from tensorflow import keras
import numpy as np
from PIL import Image
import cv2
import tensorflow as tf
print(tf.__version__)
physical_devices = tf.config.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(physical_devices[0], True)
print(physical_devices)
# %%
(train_x, train_y), (test_x, test_y) = keras.datasets.cifar100.load_data()


# %%
def block(in_x, filters, scale='down', kernel_size=3, padding='same'):
  y = keras.layers.SeparableConv2D(filters,
                                   kernel_size,
                                   padding=padding,
                                   use_bias=False)(in_x)
  y = keras.layers.BatchNormalization()(y)
  y = keras.layers.LeakyReLU()(y)
  if scale == 'down':
    y = keras.layers.MaxPooling2D()(y)
  elif scale == 'up':
    y = keras.layers.UpSampling2D()(y)
  return y


# %%
base = 128
x = keras.layers.Input((32, 32, 3))
y = keras.layers.experimental.preprocessing.Rescaling(1. / 255)(x)
y = block(y, base * 1)
y = block(y, base * 2)
y = block(y, base * 3)
y = block(y, base * 4)
y = block(y, base * 5)
y = keras.layers.Flatten()(y)
y = keras.layers.Dense(base * 5, 'relu')(y)
y = keras.layers.Reshape((1, 1, base * 5))(y)
y = block(y, base * 5, scale='up')
y = block(y, base * 4, scale='up')
y = block(y, base * 3, scale='up')
y = block(y, base * 2, scale='up')
y = block(y, base * 1, scale='up')
y = keras.layers.Conv2D(3, kernel_size=3, padding='same')(y)
y = keras.layers.experimental.preprocessing.Rescaling(255)(y)

model = keras.Model(inputs=[x], outputs=[y])
model.compile('rmsprop', 'mse')
model.summary()
model.save('auto_encoder_model.h5')

# %%
cbs = [
    keras.callbacks.EarlyStopping(patience=30, restore_best_weights=True),
    keras.callbacks.ModelCheckpoint(
        'tmp/checkpoints/auto_encoder-{epoch:03d}-{val_loss:.2f}.h5',
        save_weights_only=True,
        save_best_only=True),
    keras.callbacks.TensorBoard(log_dir='tmp/logs')
]
# %%

ds = np.concatenate([train_x, test_x])

model.fit(ds,
          ds,
          epochs=1000,
          batch_size=40,
          validation_split=0.05,
          callbacks=cbs,
          initial_epoch=194)

# %%
valid_set = test_x[:10]
res = model.predict(valid_set)
res = res.astype('int32')
from matplotlib import pyplot as plt
for i in range(10):
  plt.imshow(valid_set[i])
  plt.show()
  plt.imshow(res[i])
  plt.show()

# %%
model.load_weights('tmp/checkpoints/auto_encoder-194-318.78.h5')
# %%
model.save('auto_encoder.h5')

# %%
model = keras.models.load_model('auto_encoder_spc512_2.h5')
print(model.layers[:24])
new_m = keras.Sequential(model.layers[:24])
res = new_m.predict(valid_set)
print(res[2])

# %%
import os
img = Image.open('test.png')
x, y = img.size
img = img.crop((0, 0, max(x, y), max(x, y)))
img = img.resize((32, 32))
img = np.array(img).reshape((1, 32, 32, 3))
display(img)

res = model.predict(img)
res = res.astype('int32')
plt.imshow(img[0])
plt.show()
plt.imshow(res[0])
plt.show()

# %%
arr1 = np.random.random((10, 32, 32, 3))
arr2 = np.random.random((10, 32, 32, 3))

arr = np.concatenate((arr1, arr2))
print(arr.shape)
np.linalg.norm()